package com.jscloud.bigdata.flink.keyedstate;

import org.apache.flink.api.common.functions.FlatMapFunction;
import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.runtime.state.FunctionInitializationContext;
import org.apache.flink.runtime.state.FunctionSnapshotContext;
import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.sink.SinkFunction;
import org.apache.flink.util.Collector;

import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;

/**
 * 自定义Sink,  通过`CheckpointedFunction `接口操作 `Operator State`，实现每次达到2个元素就输出打印。
 */
public class OperatorSinkByJava {
        public static void main(String[] args) throws Exception {
                StreamExecutionEnvironment executionEnvironment = StreamExecutionEnvironment.getExecutionEnvironment();

                DataStreamSource<String> dataStreamSource = executionEnvironment.socketTextStream("bigdata01", 9999);

                SingleOutputStreamOperator<String> result = dataStreamSource.flatMap(new FlatMapFunction<String, String>() {
                        @Override
                        public void flatMap(String value, Collector<String> out) throws Exception {
                                String[] words = value.split(" ");
                                for (String word : words) {
                                        out.collect(word);
                                }
                        }
                });
                result.addSink(new MySinkFunction());
                executionEnvironment.execute();
        }
}

class MySinkFunction implements SinkFunction<String>, CheckpointedFunction {
        private int threadHold = 2;
        private List<String> list = new ArrayList<String>();

        private ListState<String> checkPointState = null;

        @Override
        public void invoke(String value, Context context) throws Exception {
                list.add(value);
                if (list.size() == threadHold) {
                        System.out.println("两个元素开始打印" + list.get(0) + "\t" + list.get(1));
                        list.clear();
                }
        }

        @Override
        public void snapshotState(FunctionSnapshotContext functionSnapshotContext) throws Exception {
                checkPointState.clear();
                for (String s : list) {
                        checkPointState.add(s);
                }

        }

        @Override
        public void initializeState(FunctionInitializationContext context) throws Exception {
                checkPointState = context.getOperatorStateStore().getListState(new ListStateDescriptor<String>("operatorState", String.class));
                if (context.isRestored()) {
                        Iterator<String> iterator = checkPointState.get().iterator();
                        while (iterator.hasNext()) {
                                list.add(iterator.next());
                        }
                }
        }
}